import struct

elfhdr = {}


def get_name(f, offset):
    name = ''
    f.seek(offset)
    while 1:
        c = f.read(1)
        if c.decode() == '\0':
            break
        else:
            name += c.decode()
    return name


def verify_elf(filename):
    f = open(filename, 'rb')
    elfident = bytearray(f.read(16))

    magic = [i for i in elfident]
    if magic[0] != 127 or magic[1] != ord('E') or magic[2] != ord('L') or magic[3] != ord('F'):
        print("your input file %s not a elf file" % filename)
        return
    else:
        temp = f.read(struct.calcsize('2HI3QI6H'))
        temp = struct.unpack('2HI3QI6H', temp)
        global elfhdr

        elfhdr['magic'] = magic
        elfhdr['e_type'] = temp[0]
        elfhdr['e_machine'] = temp[1]
        elfhdr['e_version'] = temp[2]
        elfhdr['e_entry'] = temp[3]
        elfhdr['e_phoff'] = temp[4]
        elfhdr['e_shoff'] = temp[5]
        elfhdr['e_flags'] = temp[6]
        elfhdr['e_ehsize'] = temp[7]
        elfhdr['e_phentsize'] = temp[8]
        elfhdr['e_phnum'] = temp[9]
        elfhdr['e_shentsize'] = temp[10]
        elfhdr['e_shnum'] = temp[11]
        elfhdr['e_shstrndx'] = temp[12]
    f.close()


def save_sections(elffile):
    verify_elf(elffile)
    sections = []
    global elfhdr
    sec_start = elfhdr['e_shoff']
    sec_size = elfhdr['e_shentsize']
    f = open(elffile, 'rb')
    f.seek(sec_start)
    for i in range(0, elfhdr['e_shnum']):
        temp = f.read(sec_size)
        temp = struct.unpack('2I4Q2I2Q', temp)
        sec = {'sh_name': temp[0], 'sh_type': temp[1], 'sh_flags': temp[2], 'sh_addr': temp[3], 'sh_offset': temp[4],
               'sh_size': temp[5], 'sh_link': temp[6], 'sh_info': temp[7], 'sh_addralign': temp[8],
               'sh_entsize': temp[9]}
        sections.append(sec)
    start = sections[elfhdr['e_shstrndx']]['sh_offset']
    table = []
    for i in range(0, elfhdr['e_shnum']):
        if sections[i]['sh_addr'] != 0:
            offset = start + sections[i]['sh_name']
            name = get_name(f, offset)
            filename = "mem_file/section/" + name[1:]
            table.append((name, sections[i]['sh_addr'], sections[i]['sh_offset'], sections[i]['sh_size'], filename))

    for i in range(len(table)):
        infile = open(elffile, "rb")
        infile.seek(table[i][2])
        temp = infile.read(table[i][3])
        infile.close()
        outfile = open(table[i][4], "wb")
        outfile.write(temp)
        outfile.close()
    return table, elfhdr['e_entry']


# save_sections("arch-fib")
